package dan.backend
import chisel3._
import chisel3.util._
import dan.common._ 
import dan.frontend.BrUpdateInfo

////// busy table
class BusyTableResp extends Bundle {
    val prs1_busy = Bool()
    val prs2_busy = Bool()
}

class BusyTable (
    val plWidth : Int,
    val numPregs: Int,
    val numWakeupPorts: Int
) extends CoreModule {
    val pregSize = log2Ceil(numPregs)

    val io = IO(new Bundle {
        val ren_uops = Input(Vec(plWidth, new UOp))
        val busy_resps = Output(Vec(plWidth, new BusyTableResp))
        val rebusy_reqs = Input(Vec(plWidth, Bool()))

        val wakeup_valids = Input(Vec(numWakeupPorts, Bool()))
        val wakeup_pdsts = Input(Vec(numWakeupPorts, UInt(pregSize.W)))
    })

    val busyTable = RegInit(0.U(numPregs.W))

    //将写回的寄存器置为非忙
    val busyTableWakeup = busyTable &
        ~(io.wakeup_pdsts zip io.wakeup_valids).map {
            case (pdst, valid) => valid.asUInt << pdst
        }.reduce(_|_)

    //将新分配的寄存器置为忙
    val busyTableNext = busyTableWakeup |
        (io.ren_uops zip io.rebusy_reqs) .map {
            case (uop, req) => req.asUInt << uop.physDst
        }.reduce(_|_)

    //更新busytable
    busyTable := busyTableNext

    //输出
    //这里我们将不考虑转发，转发逻辑在外面顶层模块进行
    for (i <- 0 until plWidth){
        io.busy_resps(i).prs1_busy := busyTable(io.ren_uops(i).physOprand1)
        io.busy_resps(i).prs2_busy := busyTable(io.ren_uops(i).physOprand2)
    }
}

///// freelist
class FreeList extends CoreModule {

    val io = IO(new Bundle {
        val reqs = Input(Vec(coreWidth, Bool()))
        val alloc_pregs = Vec(coreWidth, Valid(UInt(physRegBits.W)))

        // from ROB
        val dealloc_pregs = Input(Vec(coreWidth, Valid(UInt(physRegBits.W))))

        val ren_br_tags = Input(Vec(coreWidth, Valid(UInt(brTagBits.W))))

        val brupdate = Input(new BrUpdateInfo)
    })

    // all free except x0
    val freeList = RegInit(UInt(physRegNum.W), (~1.U(physRegNum.W)).asUInt)
    val allocsAfterBr = Reg(Vec(maxBrNum, UInt(physRegNum.W)))

    // allocate
    val selPregs = Wire(Vec(coreWidth, UInt(physRegNum.W)))
    val selPregsValid = VecInit(selPregs.map(_.orR))
    var mask = freeList
    for (i <- 0 until coreWidth) {
        selPregs(i) := PriorityEncoderOH(mask)
        mask = mask & (~selPregs(i)).asUInt
    }

    // maintain stored allocated regs
    val regValids = Seq.fill(coreWidth) {RegInit(false.B)}
    val regIndices = Seq.fill(coreWidth) {Reg(UInt(physRegBits.W))}
    // if selected reg is valid, always valid, otherwise may be used by req
    regValids zip selPregsValid zip io.reqs foreach { case ((regValid, selValid), req) =>
        regValid := selValid || (regValid && !req)
    }
    // whether should fill in new regs
    val selPregFire = VecInit(selPregsValid zip regValids zip io.reqs map { case ((selValid, regValid), req) =>
        (!regValid || req) && selValid
    })
    regIndices zip selPregs zip selPregFire map { case ((regNum, selPreg), fire) =>
        when (fire) {
            regNum := OHToUInt(selPreg)
        }
    }

    io.alloc_pregs zip regValids zip regIndices foreach { case ((io, valid), num) =>
        io.valid := valid
        io.bits := num
    }

    val allocOHs = regIndices map {UIntToOH(_)}
    val allocMasks = (allocOHs zip io.reqs).scanRight(0.U(physRegNum.W)) { case ((alloc, req), mask) =>
        mask | Mux(req, alloc, 0.U)
    }

    val selMask = ((selPregs zip selPregFire) map { case (reg, fire) => Mux(fire, reg, 0.U) }).reduce(_|_)
    // free pregs in mispredicted branch
    val brDeallocs = Mux(io.brupdate.resolInfo.mispredict, allocsAfterBr(io.brupdate.resolInfo.uop.brTag), 0.U)
    // pregs to free from rob
    val deallocMask =
        (io.dealloc_pregs map { de => Mux(de.valid, UIntToOH(de.bits)(physRegNum - 1, 0), 0.U) }).reduce(_|_) | brDeallocs

    val brTagValids = VecInit(io.ren_br_tags map (_.valid)).asUInt

    for (i <- 0 until maxBrNum) {
        val updateList = VecInit(io.ren_br_tags.map(_.bits === i.U)).asUInt & brTagValids
        allocsAfterBr(i) := Mux(
            updateList.orR,
            Mux1H(updateList, allocMasks.slice(1, coreWidth + 1)),
            (allocsAfterBr(i) & (~brDeallocs).asUInt) | allocMasks.head
        )
    }

    freeList := (freeList & (~selMask).asUInt | deallocMask) & (~1.U(physRegNum.W)).asUInt

}

//// maptable
class MaptableReq extends CoreBundle {
    val lrs1 = UInt(archRegBits.W)
    val lrs2 = UInt(archRegBits.W)
    val ldst = UInt(archRegBits.W)
}

class MaptableResp extends CoreBundle {
    val prs1 = UInt(physRegBits.W)
    val prs2 = UInt(physRegBits.W)
    val stale_pdst = UInt(physRegBits.W)
}

class ReMapReq extends CoreBundle {
    val ldst = UInt(archRegBits.W)
    val pdst = UInt(physRegBits.W)
    val valid = Bool()
}

class MapTable extends CoreModule {

    val io = IO(new CoreBundle{
        val map_reqs = Input(Vec(coreWidth, new MaptableReq))
        val map_resps = Output(Vec(coreWidth, new MaptableResp))

        val remap_reqs = Input(Vec(coreWidth, new ReMapReq))

        val ren_br_tags = Input(Vec(coreWidth, Valid(UInt(brTagBits.W))))

        val brupdate = Input(new BrUpdateInfo)
    })

    val mapTable = RegInit(VecInit(Seq.fill(archRegNum) { 0.U(physRegBits.W) }))
    val brShot = Reg(Vec(maxBrNum, Vec(archRegNum, UInt(physRegBits.W))))

    val remapTable = io.remap_reqs.scanLeft(mapTable) { case (table, req) =>
        VecInit(table.zipWithIndex map { case (preg, lreg) =>
            if (lreg == 0) {
                0.U
            } else {
                Mux(req.valid && req.ldst === lreg.U, req.pdst, preg)
            }
        })
    }

    // save maps at br for rollback
    io.ren_br_tags zip remapTable.slice(1, coreWidth + 1) foreach { case (tag, table) =>
        when (tag.valid) {
            brShot(tag.bits) := table
        }
    }

    mapTable := Mux(io.brupdate.resolInfo.mispredict, brShot(io.brupdate.resolInfo.uop.brTag), remapTable.last)

    // forwarding is done at top modules
    io.map_resps zip io.map_reqs foreach { case (resp, req) =>
        resp.prs1 := mapTable(req.lrs1)
        resp.prs2 := mapTable(req.lrs2)
        resp.stale_pdst := mapTable(req.ldst)
    }
}






